# -*- coding: UTF-8 -*-
"""
@Project ：swift-forge-engine 
@File    ：jwt
@IDE     ：PyCharm 
@Author  ：Tony
@Date    ：2025/2/20 11:11 
@Des     ：
"""
import json
from datetime import timedelta
from uuid import uuid4

from fastapi import Depends, Request
from fastapi.security import HTTPBearer
from fastapi.security.utils import get_authorization_scheme_param
from jose import jwt, ExpiredSignatureError, JWTError
from passlib.context import CryptContext
from pydantic_core import from_json
from sqlalchemy.ext.asyncio import AsyncSession

from App.app.admin.models.sys_user import SysUser
from App.app.admin.schemas.sys_user_schema import CurrentUserInfo
from App.common.dataclasses import AccessToken, TokenPayload, RefreshToken, NewToken
from App.common.exception.errors import AuthorizationError, ForbiddenError
from App.config.conf import settings
from App.database.mysql import async_db_session
from App.database.redis import redis_client
from App.utils.serializers import select_as_dict
from App.utils.timezone import timezone

pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

DependsJwtAuth = Depends(HTTPBearer())


def get_hash_password(password: str) -> str:
    """
    使用hash算法加密密码
    :param password: 密码
    :return: 加密后的密码
    """
    return pwd_context.hash(password)


def password_verify(plain_password: str, hashed_password: str) -> bool:
    """
    密码校验
    :param plain_password: 要验证的密码
    :param hashed_password: 要比较的hash密码
    :return: 比较密码之后的结果
    """
    return pwd_context.verify(plain_password, hashed_password)


async def create_access_token(user_id: str, multi_login: int, **kwargs) -> AccessToken:
    """
    创建TOKEN

    :param user_id: 用户ID
    :param multi_login: 是否重复登录
    :param kwargs:
    :return:
    """
    expire = timezone.now() + timedelta(seconds=settings.TOKEN_EXPIRE_SECONDS)
    session_uuid = str(uuid4())
    access_token = jwt.encode(
        {'session_uuid': session_uuid, 'exp': expire, 'sub': user_id},
        settings.TOKEN_SECRET_KEY,
        settings.TOKEN_ALGORITHM,
    )
    if multi_login == 0:  # 如果不允许重复登录，则删除该用户ID的前缀的所有TOKEN

        await redis_client.delete_prefix(f'{settings.TOKEN_REDIS_PREFIX}:{user_id}')

    await redis_client.setex(
        f'{settings.TOKEN_REDIS_PREFIX}:{user_id}:{session_uuid}',
        settings.TOKEN_EXPIRE_SECONDS,
        access_token,
    )

    # Token 附加信息单独存储
    if kwargs:
        await redis_client.setex(
            f'{settings.TOKEN_EXTRA_INFO_REDIS_PREFIX}:{session_uuid}',
            settings.TOKEN_EXPIRE_SECONDS,
            json.dumps(kwargs, ensure_ascii=False),
        )

    return AccessToken(access_token=access_token, access_token_expire_time=expire, session_uuid=session_uuid)


async def create_new_token(user_id: str, refresh_token: str, multi_login: int, **kwargs) -> NewToken:
    """
    Generate new token

    :param user_id:
    :param refresh_token:
    :param multi_login:
    :param kwargs: Access token extra information
    :return:
    """

    redis_refresh_token = await redis_client.get(f'{settings.TOKEN_REFRESH_REDIS_PREFIX}:{user_id}:{refresh_token}')
    if not redis_refresh_token or redis_refresh_token != refresh_token:
        raise ForbiddenError(message='刷新令牌已过期，请重新登录')

    new_access_token = await create_access_token(user_id, multi_login, **kwargs)
    return NewToken(
        new_access_token=new_access_token.access_token,
        new_access_token_expire_time=new_access_token.access_token_expire_time,
        session_uuid=new_access_token.session_uuid,
    )


async def create_refresh_token(user_id: str, multi_login: int) -> RefreshToken:
    """
    Generate encryption refresh token, only used to create a new token

    :param user_id: The user id of the JWT
    :param multi_login: multipoint login for user
    :return:
    """
    expire = timezone.now() + timedelta(seconds=settings.TOKEN_REFRESH_EXPIRE_SECONDS)
    refresh_token = jwt.encode(
        {'exp': expire, 'sub': user_id},
        settings.TOKEN_SECRET_KEY,
        settings.TOKEN_ALGORITHM,
    )

    if multi_login is False:
        key_prefix = f'{settings.TOKEN_REFRESH_REDIS_PREFIX}:{user_id}'
        await redis_client.delete_prefix(key_prefix)

    await redis_client.setex(
        f'{settings.TOKEN_REFRESH_REDIS_PREFIX}:{user_id}:{refresh_token}',
        settings.TOKEN_REFRESH_EXPIRE_SECONDS,
        refresh_token,
    )
    return RefreshToken(refresh_token=refresh_token, refresh_token_expire_time=expire)


def get_token(request: Request) -> str:
    """
    Get token for request header

    :return:
    """
    authorization = request.headers.get('Authorization')
    scheme, token = get_authorization_scheme_param(authorization)
    if not authorization or scheme.lower() != 'bearer':
        raise AuthorizationError(message='Token 无效')
    return token


def jwt_decode(token: str) -> TokenPayload:
    """
    Decode token

    :param token:
    :return:
    """
    try:
        payload = jwt.decode(token, settings.TOKEN_SECRET_KEY, algorithms=[settings.TOKEN_ALGORITHM])
        session_uuid = payload.get('session_uuid') or 'debug'
        user_id = payload.get('sub')
        expire_time = payload.get('exp')
        if not user_id:
            raise AuthorizationError(message='Token 无效')
    except ExpiredSignatureError:
        raise AuthorizationError(message='Token 已过期')
    except (JWTError, Exception):
        import traceback
        print(traceback.format_exc())
        raise AuthorizationError(message='Token 无效')
    return TokenPayload(id=int(user_id), session_uuid=session_uuid, expire_time=expire_time)


async def get_current_user(db: AsyncSession, pk: int) -> SysUser:
    """
    获取当前登录用户信息

    :param db:
    :param pk:
    :return:
    """

    from App.app.admin.crud.sys_user_crud import SysUserDao
    user = await SysUserDao.get_with_role(db, user_id=pk)
    if not user:
        raise AuthorizationError(message='Token 无效')
    if not user.status:
        raise AuthorizationError(message='用户已被锁定，请联系系统管理员')
    return user


def superuser_verify(request: Request) -> bool:
    """
    验证超级用户

    :param request:
    :return:
    """
    superuser = request.user.is_super
    if not superuser or not request.user.is_staff:
        raise AuthorizationError(message='非法操作')
    return superuser


async def jwt_authentication(token: str) -> CurrentUserInfo:
    """
    JWT authentication

    :param token:
    :return:
    """
    token_payload = jwt_decode(token)
    user_id = token_payload.id
    redis_token = await redis_client.get(f'{settings.TOKEN_REDIS_PREFIX}:{user_id}:{token_payload.session_uuid}')
    if not redis_token or token != redis_token:
        raise AuthorizationError(message='Token 已过期')
    cache_user = await redis_client.get(f'{settings.JWT_USER_REDIS_PREFIX}:{user_id}')
    if not cache_user:
        async with async_db_session() as db:
            current_user = await get_current_user(db, int(user_id))
            user = CurrentUserInfo(**select_as_dict(current_user))
            await redis_client.setex(
                f'{settings.JWT_USER_REDIS_PREFIX}:{user_id}',
                settings.JWT_USER_REDIS_EXPIRE_SECONDS,
                user.model_dump_json(),
            )
    else:
        user = CurrentUserInfo.model_validate(from_json(cache_user, allow_partial=True))
    return user
